from typing import TypeVar

import numpy as np
import torch

T = TypeVar("T")


def ceil_divide(x: int, y: int):
    return (x + (y - 1)) // y


def linear_trans(x0: float, x1: float, y0: float = 0.0, y1: float = 1.0, x: float | None = None):
    if x is None:
        return lambda x: (y1 - y0) / (x1 - x0) * (np.clip(x, x0, x1) - x0) + y0
    else:
        return linear_trans(x0, x1, y0, y1)(x)


def torch_gather_from_array(array: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
    """
    :param array: shape = [arr_len, *item_shape]
    :param indices: shape = [*batch_shape] dtype=int64
    :return: shape = [*batch_shape, *item_shape]
    """
    batch_shape = indices.shape
    item_shape = array.shape[1:]

    indices = torch.reshape(indices, [-1])
    values = array[indices]
    values = torch.reshape(values, [*batch_shape, *item_shape])

    return values


def torch_compile_if_possible(model: T) -> T:
    try:
        model = torch.compile(model)
    except Exception:
        pass
    return model
